"""
RSA加解密算法
"""
import math
import random
from properties import *


def encrypt(msg: bytes, e: int, n: int) -> bytes:
    """
    rsa加密函数
    """
    # 分段长度需要比n的长度小
    # 每段加密后需要进行填充, 保证解密时不会分段错误
    return crypt(msg, e, n, piece_len=n.bit_length() // 8 - 1, fill_len=FILL_LEN)


def decrypt(enc: bytes, d: int, n: int) -> bytes:
    """
    rsa解密函数
    """
    # 因为加密的时候将填充长度设置为FILL_LEN, 所以解密的时候切片需要以FILL_LEN为单位进行切片
    # fill_len设置为0代表不需要填充, 直接获取明文
    return crypt(enc, d, n, piece_len=FILL_LEN, fill_len=0)


def crypt(msg: bytes, exponent: int, n: int, piece_len: int, fill_len: int = FILL_LEN) -> bytes:
    """
    加解密函数
    :param msg: 要加解密的字节流
    :param exponent: 指数
    :param n: 模
    :param piece_len: 字节流切片长度(bytes)
    :param fill_len: 处理后的字节流填充长度(bytes), 非正数则不填充
    :return: 加解密后的字节流
    """
    # c = bytes2int(msg)
    # m = fast_pow(c, exponent, n)
    # return int2bytes(m)
    int_list = [bytes2int(msg[i:i + piece_len]) for i in range(0, len(msg), piece_len)]
    print(int_list)
    # piece_len = PIECE_LEN if len(int_list) > 1 else 0
    bytes_list = [int2bytes(fast_pow(c, exponent, n), fill_len) for c in int_list]
    return b''.join(bytes_list)


def get_rsa_keys() -> tuple[int, int, int]:
    """
    生成RSA秘钥, 包括公钥, 私钥
    :return: 公钥e, 私钥d, n
    """
    # 获取随机的p, q
    p = get_random_prime()
    print(f"p: {p}")
    q = get_random_prime()
    print(f"q: {q}")
    while p == q:
        p = get_random_prime()

    # 计算n和phi(n)
    n = p * q
    phi_n = (p - 1) * (q - 1)

    # 计算e
    e = 1 if RANDOM_EXPONENT else DEFAULT_EXPONENT
    while True:
        if phi_n % e != 0:  # === if math.gcd(e, phi_n) == 1: because e is a prime
            break
        e = get_random_prime(DEFAULT_EXPONENT.bit_length())

    # 计算d
    d = inverse(e, phi_n)
    print(f"n: {n}")
    print(f"e: {e}")
    print(f"d: {d}")
    print(f"phi_n: {phi_n}")
    return e, d, n


def inverse(a: int, n: int) -> int:
    """
    使用拓展欧几里得算法求模的逆元
    返回 1 / a (mod n).
    @a and @n 必须互质
    """
    g, inv, _ = extended_gcd(a, n)
    if g != 1:
        raise ValueError(f"不存在逆元: a={a}, n={n}")
    if (inv * a) % n != 1:
        raise ValueError(f"逆元计算错误")
    return inv


def extended_gcd(a: int, b: int) -> tuple[int, int, int]:
    """
    拓展欧几里得算法
    返回(g, x, y) : gcd(a, b) = g = ax + by
    """
    x = 0
    y = 1
    lx = 1
    ly = 0
    oa = a
    ob = b
    while b != 0:
        q = a // b
        a, b = b, a % b
        x, lx = (lx - (q * x)), x
        y, ly = (ly - (q * y)), y
    # 只返回正数结果
    if lx < 0:
        lx += ob
    if ly < 0:
        ly += oa
    return a, lx, ly


def get_random_prime(size=PRIME_LEN) -> int:
    """
    生成指定位数的随机素数
    :param size: 生成素数的位数
    :return: 生成的随机素数
    """
    while True:
        x = get_random_odd_int(size=size)
        if is_prime(x):
            return x


def get_random_odd_int(size: int) -> int:
    """
    生成对应位数的, 首位均为1的随机数
    :param size: 随机数位数
    :return: 随机数
    """
    return random.getrandbits(size) | (1 << (size - 1)) | 1


def is_prime(x: int) -> bool:
    """
    检测一个数字是否为素数
    :param x: 要检测的数字
    :return: 是否大概率为素数(>99%)
    """
    # 利用python的小整数池, 提前筛选一部分, 提高效率
    smallPrimes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101,
                   103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199,
                   211, 223, 227, 229, 233, 239, 241, 251]
    if x in smallPrimes:
        return True
    for i in smallPrimes:
        if x % i == 0:
            return False
    return miller_rabin(x)


def miller_rabin(n: int, nRound: int = 5) -> bool:
    """
    MillerRabin算法, 简单检验是否可能为素数
    :param n: 需要检测的数字
    :param nRound:
    :return: 是否应该是素数(75%)
    """
    # 防止n=1时进入死循环
    if n < 2:
        return False

    # (2 ^ k) * q = (n - 1)
    k, q = 0, n - 1
    while (q & 1) == 0:
        k += 1
        q >>= 1

    for _ in range(nRound):
        a = random.randrange(2, n - 1)
        x = fast_pow(a, q, n)
        if x == 1 or x == n - 1:
            continue
        for _ in range(k - 1):
            x = fast_pow(x, 2, n)
            # 以后永远都是1, 没有必要再算了
            if x == 1:
                return False
            if x == n - 1:
                break
        else:
            return False
    return True


def fast_pow(base: int, exp: int, mod: int) -> int:
    """
    快速幂算法
    """
    # pow(base, exp, mod)
    y = 1
    while True:
        if exp == 0:
            return y
        while exp > 0 and exp % 2 == 0:
            base = (base * base) % mod
            exp //= 2
        exp -= 1
        y = (base * y) % mod


def bytes2int(raw_bytes: bytes) -> int:
    """
    将字节流转化为int
    """
    return int.from_bytes(raw_bytes, "big", signed=False)


def int2bytes(num: int, fill_len: int = 0) -> bytes:
    """
    将int转化为字节流
    """
    if fill_len > 0:
        return num.to_bytes(fill_len, 'big')
    bytes_len = max(1, math.ceil(num.bit_length() / 8))
    return num.to_bytes(bytes_len, 'big')


def main():
    msg = "2002 A.M. TURING AWARD. RSA, an acronym for Rivest, Shamir and Adleman, uses algorithmic number theory to provide an efficient realization of a public-key cryptosystem, a concept first envisioned theoretically by Whitfield Diffie, Martin Hellman and Ralph Merkle. RSA is now the most widely used encryption method, with applications throughout the Internet to secure on-line transactions. It has also inspired breakthrough work in both theoretical computer science and mathematics."
    # msg = 'CsCN'*9
    # print(len(msg), PIECE_LEN)
    print(f"KEY_LEN: {KEY_LEN}")
    e, d, n = get_rsa_keys()
    enc = encrypt(msg.encode(ENCODING), e=e, n=n)
    print(f"加密后的字节流: {enc}")
    dec = decrypt(enc, d, n)
    print(f"解密后的文本: {dec.decode(ENCODING)}")


def test():
    global KEY_LEN
    KEY_LEN = 256
    for _ in range(1000):
        main()
    print('OK')
    exit()


if __name__ == '__main__':
    # test()
    main()
